#include "trt_run.h"

#include <fstream>
#include <cuda_runtime_api.h>
#include "NvOnnxParser.h"

class MyLogger : public nvinfer1::ILogger {
public:
    explicit MyLogger(nvinfer1::ILogger::Severity severity = nvinfer1::ILogger::Severity::kWARNING) : severity_(
            severity) {}

    void log(nvinfer1::ILogger::Severity severity, const char *msg) noexcept override {
        if (severity <= severity_) {
            std::cerr << msg << std::endl;
        }
    }

    nvinfer1::ILogger::Severity severity_;
};

// 返回每种数据类型占用的字节数
inline unsigned int getElementSize(nvinfer1::DataType t) {
    switch (t) {
        case nvinfer1::DataType::kINT32:
            return 4;
        case nvinfer1::DataType::kFLOAT:
            return 4;
        case nvinfer1::DataType::kHALF:
            return 2;
        case nvinfer1::DataType::kINT8:
            return 1;
        default:
            throw std::runtime_error("Invalid DataType.");
            return 0;
    }
}


YoloDetect::YoloDetect(const std::string &engine_file, const float &conf_threshold,
                       const std::vector<std::string> &class_names, const bool &draw_box) {
    this->confthreshold = conf_threshold;
    this->class_names = class_names;
    int num_classes = (int) this->class_names.size();
    this->draw_box = draw_box;
    // 初始化画框颜色
    if (this->draw_box) {
        this->SetColor(num_classes);
    }

    // 设置显卡编号
    cudaSetDevice(0);
    // 加载trt引擎
    std::string cached_engine;
    // 创建读写流对象
    std::fstream file;
    file.open(engine_file, std::ios::binary | std::ios::in);

    // 文件当前指针不为-1(EOF为末尾标识符值为-1)
    while (file.peek() != EOF) {
        std::stringstream buffer;
        // 把file中的流对象放入buffer中
        buffer << file.rdbuf();
        // 转为字符串添加到cached_engine中
        cached_engine.append(buffer.str());
    }
    file.close();

    // 创建模型实例
    MyLogger logger;
    nvinfer1::IRuntime *trtRuntime = nvinfer1::createInferRuntime(logger);

    // 模型读入缓冲区，反序列化引擎
    engine = trtRuntime->deserializeCudaEngine(cached_engine.data(), cached_engine.size(), nullptr);

    // 存储中间值的空间
    context = engine->createExecutionContext();

    // bingding是缓冲区获取参数，Tensor是从张量获取参数
    // 获取输入和输出的总数  2
    int nbTensors = engine->getNbIOTensors();
    buffer_size.resize(nbTensors);

    // 获取输入输出名称
    if (nbTensors == 2) {
        input_names = engine->getIOTensorName(0);
        output_names = engine->getIOTensorName(1);
    } else {
        throw std::runtime_error("输入输出总个数与接收值不匹配.");
        exit(0);
    }

    // 获取输入维度和尺寸 {1,3,640,640} [n,c,h,w]
    nvinfer1::Dims input_dim = engine->getTensorShape(input_names);
    // 输入数据类型
    nvinfer1::DataType input_dtype = engine->getTensorDataType(input_names);
    // 输入网络图像的宽高
    this->input_size2f = cv::Size2f((float) input_dim.d[3], (float) input_dim.d[2]);
    // 计算输入流的总形状大小
    int input_size = 1;
    for (int j = 0; j < input_dim.nbDims; ++j) {
        input_size *= input_dim.d[j];
    }
//    std::cout << getElementSize(input_dtype) << std::endl;  // 4
    // 计算输入的总空间大小
    auto input_sizeof = input_size * getElementSize(input_dtype);
    buffer_size[0] = input_sizeof;
    cudaMalloc(&buffers[0], input_sizeof);

    // 获取输出尺寸 {1,84,8400} [n, (cx,cy,w,h,类别置信度), (20*20+40*40+80*80)]
    nvinfer1::Dims output_dim = engine->getTensorShape(output_names);
    nvinfer1::DataType output_dtype = engine->getTensorDataType(output_names);
    // 计算输出总形状
    for (int j = 0; j < output_dim.nbDims; ++j) {
        output_shape.push_back(output_dim.d[j]);
        output_size *= output_dim.d[j];
    }
    // 计算输出占用内存总大小
    auto output_sizeof = output_size * getElementSize(output_dtype);
    buffer_size[1] = output_sizeof;
    cudaMalloc(&buffers[1], output_sizeof);

    //创建cuda流
    cudaStreamCreate(&stream);
}

YoloDetect::~YoloDetect() {
    // 释放流和缓冲区
    cudaStreamDestroy(stream);
    cudaFree(buffers[0]);
    cudaFree(buffers[1]);

    // 破坏引擎
    context->destroy();
    engine->destroy();
}

std::vector<Detection> YoloDetect::detect(cv::Mat &draw_img) {
    // 网络前处理时间: 3ms
    // resize的图像和归一化之后的图像
    cv::Mat resize_image, float_image;
    // 存储缩放比例，左边和上边 pad 的值
    std::map<std::string, float> m_resize;
    // 真正输入网络的形状
//    std::vector<int64_t> inputnew_shape = this->inputtensor_shape;  // {1,3,640,640}
    cv::Size realimg_shape = draw_img.size();  // 1080 x 810

    // resize等比缩放  [1080 x 810 --> 640 x 640]
    this->ScaleResize(draw_img, resize_image, m_resize);

    // BGR --> RGB 标准化
    this->ImgNormalize(resize_image, float_image);

    // 堆区创建一个数组，大小519168
    auto *blob = new float[float_image.cols * float_image.rows * float_image.channels()];

    // h,w,c --> c,h,w
    YoloDetect::ImgHwc2chw(float_image, blob);

    // 网络推理时间: 15ms
    // 分配到cuda输入数据
    cudaMemcpyAsync(buffers[0], blob, buffer_size[0], cudaMemcpyHostToDevice, stream);

    // 进行异步推理
    context->enqueueV2(buffers, stream, nullptr);

    // cuda输出结果拷贝出来到cpu得到输出数据
    auto *output_tensor = new float[output_size];
    cudaMemcpyAsync(output_tensor, buffers[1], buffer_size[1], cudaMemcpyDeviceToHost, stream);
    cudaStreamSynchronize(stream);

    double start1 = clock();

    // (1, 84, 8400)-->(1, 8400, 84)
    int output_channel = (int) output_shape[0];  // 1
    int output_row = (int) output_shape[1];  // 84
    int output_col = (int) output_shape[2];  // 8400
    // 多张图片这里要循环rows然后reshape，直接reshape不对，会把相邻几个值作为通道的值，和np的reshape不一样
    cv::Mat output_mat(output_channel, output_row * output_col, CV_32F, output_tensor);

    std::vector<cv::Mat> output_mat_vec;
    for (int i = 0; i < output_channel; i++) {
        cv::Mat output_one = output_mat.row(i);
        // cn为0表示通道不变，1表示一个通道。行是output_channel，列自动计算。转置交换后两个坐标轴
        cv::Mat output_transpose_one = output_one.reshape(1, output_row).t();
        output_mat_vec.push_back(output_transpose_one);
    }
    // 取第一张图片的结果，有batch需要改
    cv::Mat output_transpose = output_mat_vec[0];
    auto* output_tensor_transpose = output_transpose.ptr<float>();

    double end1 = clock();
    double t1 = (end1 - start1) / CLOCKS_PER_SEC * 1000;  // detect_time: 2ms
    std::cout << "detect_time_transpose: " << t1 << std::endl;

    // 输出处理：置信度阈值筛选，nms，坐标反算 2ms
    std::vector<Detection> result = this->ScaleNms(output_tensor_transpose, realimg_shape, m_resize);

    delete[] blob;
    delete[] output_tensor;

    // 画框
    if (this->draw_box) {
        this->DrawBox(draw_img, result);
    }

    return result;
}

// 画框
void YoloDetect::DrawBox(cv::Mat &draw_img, const std::vector<Detection> &detections) {
    for (const Detection &detection: detections) {
        int clas_index = detection.clas_id;
        int conf = (int) std::round(detection.conf * 100);
        int x = detection.box.x;
        int y = detection.box.y;
//        int h = detection.box.height;

        std::string label = class_names[clas_index] + ": 0." + std::to_string(conf);

        cv::Scalar color = this->clas_colors[clas_index];
        // 画框
        cv::rectangle(draw_img, detection.box, color, 2);
        // 写字
        cv::putText(draw_img, label, cv::Point(x, y - 8), cv::LINE_AA, 0.8, color, 2);
    }
}

void YoloDetect::SetColor(const int &num_classes) {
    int size = 3;  // 最多有 3^3=27 种颜色
    int *arr_b = new int[size]{0, 255, 128};
    int *arr_g = new int[size]{255, 128, 0};
    int *arr_r = new int[size]{128, 0, 255};

    cv::Scalar color;
    // 循环类别总数
    for (int i = 0; i < num_classes; i++) {
        // 转成5进制数
        int c = i % size;
        int ab = (int) i / size;
        int b = ab % 10;
        int a = (int) ab / 10;

        color = cv::Scalar(arr_b[a], arr_g[b], arr_r[c]);

        this->clas_colors.emplace_back(color);
    }

    delete[] arr_b;
    delete[] arr_g;
    delete[] arr_r;
}


// 缩放图像加黑边
void YoloDetect::ScaleResize(const cv::Mat &image, cv::Mat &resize_image, std::map<std::string, float> &m_resize) {
    float h = (float) image.size().height;
    float w = (float) image.size().width;
    int dims = image.channels();
    float IMG_HEIGHT = this->input_size2f.height;
    float IMG_WIDTH = this->input_size2f.width;

    float scale = std::min(IMG_HEIGHT / h, IMG_WIDTH / w);  // 缩放比例
    m_resize.insert(std::pair<std::string, float>("scale", scale));

    cv::resize(image, resize_image, cv::Size(0, 0), scale, scale, cv::INTER_CUBIC);  // 等比缩放

    float r_h = (float) resize_image.size().height;  // 变形后的大小
    float r_w = (float) resize_image.size().width;

    int top = int((IMG_HEIGHT - r_h) / 2);
    int bottom = int(IMG_HEIGHT - r_h - (float) top);
    int left = int((IMG_WIDTH - r_w) / 2);
    int right = int(IMG_WIDTH - r_w - (float) left);

    m_resize.insert(std::pair<std::string, float>("top", (float) top));
    m_resize.insert(std::pair<std::string, float>("left", (float) left));

    copyMakeBorder(resize_image, resize_image, top, bottom, left, right, cv::BORDER_CONSTANT,
                   cv::Scalar(114, 114, 114));
}

// BGR --> RGB 标准化 除255
void YoloDetect::ImgNormalize(const cv::Mat &resize_image, cv::Mat &float_image) {
    cv::cvtColor(resize_image, float_image, cv::COLOR_BGR2RGB);

    // 改变数据类型，改成float，不改变通道数
    float_image.convertTo(float_image, CV_32FC3, 1.0 / 255);

}

// BGR --> RGB 标准化 减均值除方差
void YoloDetect::ImgNormalize_(const cv::Mat &resize_image, cv::Mat &float_image) {
    // 改成float
    resize_image.convertTo(resize_image, CV_32FC3);

    //数据处理 标准化
    std::vector<cv::Mat> channels, channel_p;
    cv::split(resize_image, channels);
    cv::Mat R, G, B;
    B = channels.at(0);
    G = channels.at(1);
    R = channels.at(2);

    B = (B / 255. - 0.406) / 0.225;
    G = (G / 255. - 0.456) / 0.224;
    R = (R / 255. - 0.485) / 0.229;

    channel_p.push_back(R);
    channel_p.push_back(G);
    channel_p.push_back(B);

    cv::merge(channel_p, float_image);
}

void YoloDetect::ImgHwc2chw(const cv::Mat &float_image, float *blob) {
    // cols->w,rows->h
    cv::Size floatimage_size{float_image.cols, float_image.rows};

    // h,w,c --> c,h,w;
    std::vector<cv::Mat> chw(float_image.channels());
    for (int i = 0; i < float_image.channels(); ++i) {
        // blob是数组首地址
        // 数组指针地址变形为 floatImageSize大小，成为 Mat 对象，赋给容器的第一个元素。数组地址加一个通道(w*h)的大小，进行下一次赋值;
        chw[i] = cv::Mat(floatimage_size, CV_32FC1, blob + i * floatimage_size.width * floatimage_size.height);
    }
    // 用split将图像切分通道赋值给chw容器，同时也是赋值给blob数组
    cv::split(float_image, chw);
}

// 网络输出值地址，检测图片形状，缩放比例和pad
std::vector<Detection> YoloDetect::ScaleNms(const float *output_tensor, const cv::Size &realimg_shape,
                                            const std::map<std::string, float> &m_resize) {
    std::vector<cv::Rect> boxes;  // 框
    std::vector<float> confs;  // 置信度
    std::vector<int> clas_ids;  // 类别索引

    // 输出形状  [1,8400,84]  [x1,y1,x2,y2，置信度，类别索引]
    // 输出大小 1*8400*84=705600
    size_t count = output_size;
    std::vector<float> output(output_tensor, output_tensor + count);

    // 模型类别数  output_shape:[1,84,8400]
    int num_classes = (int) output_shape[1] - 4;
    // 单个批次的数据大小
    int elementsInBatch = (int) (output_shape[1] * output_shape[2]);

    // 获取缩放比例和左边上边填充像素值
    float scale = (float) m_resize.find("scale")->second;
    float top_pad = (float) m_resize.find("top")->second;
    float left_pad = (float) m_resize.find("left")->second;

    // 循环每个输出的框直到循环完所有框
    for (auto it = output.begin(); it != output.begin() + elementsInBatch; it += output_shape[1]) {
        int clas_id = int(std::max_element(it + 4, it + 4 + num_classes) - (it + 4));  //std::max_element返回范围内的最大元素
        float confidence = it[4 + clas_id];

        // 框的置信度小于阈值直接跳过
        if (confidence < this->confthreshold) {
            continue;
        }

        // 获取cx,cy,w,h
        float cx = it[0];
        float cy = it[1];
        float ow = it[2];
        float oh = it[3];

        // 坐标反算，计算box
        cv::Rect box;
        // std::round 四舍五入
        int x = (int) std::round(((cx - ow / 2) - left_pad) / scale);
        int y = (int) std::round(((cy - oh / 2) - top_pad) / scale);
        int width = (int) std::round(ow / scale);
        int height = (int) std::round(oh / scale);

        // 把box的坐标限制在图像内
        box.x = std::max(0, x);
        box.y = std::max(0, y);
        int x2 = std::min(x + width, realimg_shape.width);
        int y2 = std::min(y + height, realimg_shape.height);
        box.width = x2 - box.x;
        box.height = y2 - box.y;

        // emplace_back 和 push_back 比省略了拷贝和移动元素的过程
        boxes.emplace_back(box);
        confs.emplace_back(confidence);
        clas_ids.emplace_back(clas_id);
    }

    // NMS获取符合条件的索引
    std::vector<int> indexes;
    // iou_threshold: 0.45
    cv::dnn::NMSBoxes(boxes, confs, this->confthreshold, 0.45, indexes);

    // 创建结构体，结果存放在容器中
    std::vector<Detection> detections;
    for (int idx: indexes) {
        Detection det;
        det.box = cv::Rect(boxes[idx]);
        det.conf = confs[idx];
        det.clas_id = clas_ids[idx];
        detections.emplace_back(det);
    }
    return detections;
}
